#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <errno.h>
#include <stddef.h>
#include <unistd.h>
#include "nvml.h"

#define MAX_CONNECT_NUM 5
#define BUFFER_SIZE 8192
#define PIDFile "/var/run/nvstatd.pid"

const char *filename = "/var/run/nvstat.sock";

int write_pid_file(const char *fname);

int get_utility(char* buffer, int size) {
        nvmlUtilization_t util;
        nvmlDevice_t device;
        unsigned int deviceCount;
        int i = 0;
        char uuid[64];
        int ret;
        char *p = buffer;
        int n = size;

        ret = nvmlDeviceGetCount_v2(&deviceCount);
        if ( ret != NVML_SUCCESS ) {
                perror(nvmlErrorString(ret));
                return -1;
        }
#ifdef DEBUG
        perror("nvmlDeviceGetCount_v2");
#endif

        ret = snprintf(p, n, "{");
        p++;
        n--;

        for (i = 0; i < deviceCount; i++) {
                ret = nvmlDeviceGetHandleByIndex_v2 (i, &device);
                if ( ret != NVML_SUCCESS ) {
                        perror(nvmlErrorString(ret));
                        continue;
                }
#ifdef DEBUG
                perror("nvmlDeviceGetHandleByIndex_v2");
#endif
                ret = nvmlDeviceGetUUID(device, uuid, 63);
                if ( ret != NVML_SUCCESS ) {
                        perror(nvmlErrorString(ret));
                        continue;
                }
#ifdef DEBUG
                perror("nvmlDeviceGetUUID");
#endif
                ret = nvmlDeviceGetUtilizationRates(device, &util);
                if ( ret != NVML_SUCCESS ) {
                        perror(nvmlErrorString(ret));
                        continue;
                }
#ifdef DEBUG
                perror("nvmlDeviceGetUtilizationRates");
#endif
                ret = snprintf(p, n, "\"%s\":{\"gpu_util\":%d,\"memory_util\":%d,\"encoder_util\":-1,\"decoder_util\":-1,\"tx_util\":-1,\"rx_util\":-1}%c", 
                        uuid, 
                        util.gpu, 
                        util.memory, 
                        ',');
                if ( ret < 0) {
                        *(p+n) = '\0';
                        goto ret1;
                }
                printf("n=%d,ret=%d\n",n,ret);
                p += ret > n ? n : ret;
                n -= ret > n ? n : ret;
            }
            if( *(p-1) == ',' ) {
                    p--; // eat ','
                    snprintf(p++, n+1, "}");
            }
ret1:
            return p-buffer;
}


int do_server() {
        int fd, new_fd, len, i;
        struct sockaddr_un un;
        int ret;
        int n , left;
        char *p;

        fd = socket(AF_UNIX, SOCK_STREAM, 0);
        if(fd < 0){
                perror("Request socket failed.");
                exit(-1);
        }
        
        unlink(filename);
        un.sun_family = AF_UNIX;
        strcpy(un.sun_path, filename);
        if( bind( fd, (struct sockaddr *)&un, sizeof(un) ) < 0 ) {
                perror("bind failed.");
                exit(-1);
        }

        if( listen( fd, MAX_CONNECT_NUM ) < 0 ) {
                perror("listen failed.");
                exit(-1);
        }
        
        // allow normal user access
        chmod(filename, 0777);

        while(1){
                struct sockaddr_un client_addr;
                char buffer[BUFFER_SIZE];

                bzero(buffer, BUFFER_SIZE);
                len = sizeof(client_addr);

                new_fd = accept(fd, NULL, NULL);
                if(new_fd < 0){
                        perror("accept failed\n");
                        continue;
                }
                ret = get_utility(buffer, BUFFER_SIZE);
              
                printf("send \"%s\", len %d\n", buffer, ret+1);
                n = left = ret + 1;
                p = buffer;
                while( left > 0) {
                     n  = send( new_fd, p, left, 0 );
                     if( n > 0 ) {
                             printf("%d send\n", n);
                             p += n;
                             left -= n;
                     }
                }
                close(new_fd);
        }
        close(fd);
}

//#define DEBUG

/**
 *  https://docs.nvidia.com/deploy/nvml-api/
 *
 *
 */
int main(int argc, char* argv[])
{
        char buffer[1024];

        int r = daemon(0,0);
        if ( r != 0 ) {
                perror(strerror(errno));
                exit(r);
        }

        r = write_pid_file(PIDFile);
        if ( r != 0 ) {
                perror(strerror(errno));
                exit(r);
        }

        nvmlReturn_t ret = nvmlInitWithFlags(0);
        if ( ret != NVML_SUCCESS ) {
                perror(nvmlErrorString(ret));
                return ret;
        }
#ifdef DEBUG
        perror("nvmlInitWithFlags");
#endif

        ret = get_utility(buffer, sizeof(buffer));

        printf("%s\n%d\n", buffer, ret);

        do_server(buffer, sizeof(buffer));

err1:
        nvmlShutdown();
        return ret;
}

